import requests
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_absolute_error, r2_score
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Bidirectional, GRU
import tushare as ts
import math
import random
import webbrowser

# 设置随机种子，保证结果的可复现性
random.seed(1)
np.random.seed(1)
tf.random.set_seed(1)

# 获取股票数据的函数，使用Tushare的数据接口
def get_stock_data(token, ts_code, start_date, end_date, fields):
    try:
        pro = ts.pro_api(token)  # 请替换为自己的Tushare token
        data = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date, fields=fields)
        return data
    except Exception as e:
        print(f"获取股票数据时出现错误: {e}")
        return None

# 数据预处理函数，进行缺失值处理、数据归一化等操作
def preprocess_data(data):
    df = data.copy()
    df.set_index('trade_date', inplace=True)
    df.index = pd.to_datetime(df.index)
    df.dropna(inplace=True)
    columns = ['open', 'high', 'low', 'close', 'vol']
    df = df[columns]
    scaler = MinMaxScaler(feature_range=(0, 1))
    df_scaled = scaler.fit_transform(df)
    return df, df_scaled, scaler

# 创建训练数据集和测试数据集的函数
def create_dataset(dataset, time_steps=1):
    dataX, dataY = [], []
    for i in range(len(dataset) - time_steps - 1):
        a = dataset[i:(i + time_steps), :]
        dataX.append(a)
        dataY.append(dataset[i + time_steps, -1])  # 预测收盘价
    return np.array(dataX), np.array(dataY)

# 构建并训练LSTM模型的函数，支持不同模型类型选择
def build_and_train_model(model_type, X_train, y_train, X_test, y_test):
    model = Sequential()
    if model_type == 1:
        model.add(LSTM(units=100, activation='relu', input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(Dense(units=1))
    elif model_type == 2:
        model.add(LSTM(units=100, activation='relu', return_sequences=True, input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(LSTM(units=100, activation='relu'))
        model.add(Dense(1))
    elif model_type == 3:
        model.add(Bidirectional(LSTM(100, activation='relu'), input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(Dense(1))
    elif model_type == 4:
        model.add(GRU(units=100, activation='relu', input_shape=(X_train.shape[1], X_train.shape[2])))
        model.add(Dense(1))
    model.compile(loss='mean_squared_error', optimizer='adam')
    history = model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=50, batch_size=64, verbose=1)
    return model, history

# 计算均方根误差（RMSE）和平均绝对误差（MAE）
def calculate_metrics(y_true, y_pred):
    rmse = math.sqrt(((y_true - y_pred) ** 2).mean())
    mae = mean_absolute_error(y_true, y_pred)
    r2 = r2_score(y_true, y_pred)
    return rmse, mae, r2

if __name__ == "__main__":
    token = '572b6125c826e36619a670a9943995412179dabcfb739c5c48913f46'
    ts_code = "600519.SH"  # 股票代码
    start_date = "20100101"
    end_date = "20241225"
    fields = ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    # 获取股票数据
    stock_data = get_stock_data(token, ts_code, start_date, end_date, fields)
    if stock_data is None or stock_data.empty:
        print("没有获取到有效的股票数据！")
        exit(1)

    # 数据预处理
    df, df_scaled, scaler = preprocess_data(stock_data)

    # 划分训练集和测试集
    train_size = int(len(df_scaled) * 0.8)
    train_data, test_data = df_scaled[0:train_size, :], df_scaled[train_size:len(df_scaled), :]
    time_steps = 30  # 增加时间步长
    X_train, y_train = create_dataset(train_data, time_steps)
    X_test, y_test = create_dataset(test_data, time_steps)

    # 调整输入形状
    X_train = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], df.shape[1]))
    X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], df.shape[1]))

    # 选择模型类型（1: 单层 LSTM; 2: 多层 LSTM; 3: 双向 LSTM; 4: GRU）
    model_type = 2  # 尝试多层 LSTM
    model, history = build_and_train_model(model_type, X_train, y_train, X_test, y_test)

    # 预测训练集和测试集
    train_predict = model.predict(X_train)
    test_predict = model.predict(X_test)

    # 反归一化
    train_predict = scaler.inverse_transform(np.concatenate((train_data[time_steps:], train_predict), axis=1))[:, -1]
    test_predict = scaler.inverse_transform(np.concatenate((test_data[time_steps:], test_predict), axis=1))[:, -1]

    # 计算RMSE, MAE, R2
    test_rmse, test_mae, test_r2 = calculate_metrics(df['close'].iloc[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1].values, test_predict)
    print(f"测试集RMSE: {test_rmse}")
    print(f"测试集MAE: {test_mae}")
    print(f"测试集R2: {test_r2}")

    # 绘制股票价格走势与预测结果对比图
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df.index[time_steps:len(train_predict) + time_steps - 1], y=df['close'].iloc[time_steps:len(train_predict) + time_steps - 1], mode='lines', name='真实股价（训练集）'))
    fig.add_trace(go.Scatter(x=df.index[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1], y=df['close'].iloc[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1], mode='lines', name='真实股价（测试集）'))
    fig.add_trace(go.Scatter(x=df.index[time_steps:len(train_predict) + time_steps - 1], y=train_predict, mode='lines', name='预测股价（训练集）', line=dict(dash='dash')))
    fig.add_trace(go.Scatter(x=df.index[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1], y=test_predict, mode='lines', name='预测股价（测试集）', line=dict(dash='dash')))
    fig.update_layout(title='股票真实价格与预测价格对比', xaxis_title='日期', yaxis_title='股价')
    fig.show()
